import os
import re
import logging
from threading import Thread, Event
import paramiko
from cloudcontrol.exceptions import AuthenticationError, NonzeroReturnStatusError
from cloudcontrol._support import bytes_from_stream
from cloudcontrol._support.pickledexec import PickledExecution
from cloudcontrol._connections.base import AbstractConnection
from cloudcontrol._logging import log_event
from cloudcontrol._streams import CommandStream, GeneratorStream

_secret_unittest_ssh_key_filename = None
"""SPI used by the unit testing for the sshd loopback"""


class _SFTPFileAdapter(object):
    """helper object that gives SFTP file objects from Paramiko
    a 'name' property (which is missing for some reason)"""
    
    def __init__(self, sftp_file, filepath):
        self.__sftp_file = sftp_file
        self.__filepath = filepath
    
    def __getattr__(self, attr_name):
        if attr_name == "name":
            return self.__filepath
        return getattr(self.__sftp_file, attr_name)


class RemoteConnection(AbstractConnection):
    """connection object that correponds to a remote SSH connection"""
    
    def __init__(self, username, hostname, getpass, port):
        AbstractConnection.__init__(self)
        
        # keep track of our connection metadata
        self.__username = username
        self.__hostname = hostname
        self.__port = int(port)
        self.__getpass_callback = getpass
        
        # create our internal ssh client
        self.__ssh_client = paramiko.SSHClient()
        self.__ssh_client.load_system_host_keys()
        # TODO: use WarningPolicy instead?
        # self.__ssh_client.set_missing_host_key_policy( paramiko.WarningPolicy() )
        self.__ssh_client.set_missing_host_key_policy( paramiko.AutoAddPolicy() )
        self.__sftp_client = None
        
        # allow explicit setting of ssh keys
        # TODO: expose this as API for the connections and users modules
        if _secret_unittest_ssh_key_filename:
            self.__key_filenames = [_secret_unittest_ssh_key_filename]
        else:
            self.__key_filenames = []
        
        # set up logging
        self.__logger = logging.getLogger("cloudcontrol.connections")
    
    def __del__(self):
        self.__ssh_client.close()
        if self.__sftp_client:
            self.__sftp_client.close()
    
    def _ensure_sftp_is_ready(self):
        if not self.__sftp_client:
            log_event(self.__logger.info, "connecting_to_sftp")
            self._ensure_connected_and_authenticated()
            try:
                self.__sftp_client = self.__ssh_client.open_sftp()
            except Exception, e:
                log_event(self.__logger.info, "sftp_connection_failed")
                self.__sftp_client.close()
                raise
    
    def _ensure_connected_and_authenticated(self):
        log_event(self.__logger.info, "connecting_to_ssh")
        
        # keep track of the last exception we got from paramiko during authentication attempts
        ssh_failure_reasons = []
        
        # only do this if we are not authenticated yet
        def is_connected_and_authenticated():
            return self.__ssh_client.get_transport() and self.__ssh_client.get_transport().is_authenticated()
        
        # ATTEMPT 1: use local PKeys
        if not is_connected_and_authenticated():
            log_event(self.__logger.info, "trying_ssh_auth_via_keys")
            
            try:
                if len(self.__key_filenames) > 0:
                    
                    # we have an explicit list of ssh keys to try
                    log_event(self.__logger.info, "using_explicit_ssh_keys_list")
                    self.__ssh_client.connect(
                        hostname = self.__hostname,
                        username = self.__username,
                        port = self.__port,
                        key_filename = self.__key_filenames
                    )
                    
                else:
                    
                    # use the default keys
                    log_event(self.__logger.info, "using_default_ssh_keys_list")
                    self.__ssh_client.connect(
                        hostname = self.__hostname,
                        username = self.__username,
                        port = self.__port,
                    )
            except paramiko.SSHException, e:
                ssh_failure_reasons.append(str(e))
                self.__ssh_client.close()
            
        # ATTEMPT 2: use ssh agent on the local machine
        if not is_connected_and_authenticated():
            log_event(self.__logger.info, "trying_ssh_auth_via_sshagent")
        
            try:
                agent = paramiko.Agent()
                
                for pkey in agent.get_keys():
                    try:
                        self.__ssh_client.connect(
                            hostname = self.__hostname,
                            username = self.__username,
                            port = self.__port,
                            pkey = pkey,
                            allow_agent = True,
                            look_for_keys = True
                        )
                    except paramiko.SSHException, e:
                        ssh_failure_reasons.append(str(e))
                        self.__ssh_client.close()
                        agent.close()
                        
            except paramiko.SSHException, e:
                ssh_failure_reasons.append(str(e))
                self.__ssh_client.close()
                agent.close()
            
        # ATTEMPT 3: final attempt, use password
        if not is_connected_and_authenticated():
            
            try:
                if self.__getpass_callback:
                    log_event(self.__logger.info, "trying_ssh_auth_via_password")
                    password = self.__getpass_callback()
                    self.__ssh_client.connect(
                        hostname = self.__hostname,
                        username = self.__username,
                        port = self.__port,
                        password = password
                    )
                else:
                    # failure, we didn't even have that
                    log_event(self.__logger.warn, "missing_ssh_password_callback")
                    self.__ssh_client.close()
                    
            except paramiko.SSHException, e:
                ssh_failure_reasons.append(str(e))
                self.__ssh_client.close()
        
        # final check
        if not is_connected_and_authenticated() and len(ssh_failure_reasons) > 0:
            # re-raise if we *still* aren't authenticated
            log_event(self.__logger.warn, "failed_ssh_auth")
            raise AuthenticationError("failed to authenticate for '%(username)s@%(hostname)s' after trying accessible ssh keys, sshagent, and password callback. (reasons: %(paramiko_reasons)s)" % dict(
                username = self.__username,
                hostname = self.__hostname,
                paramiko_reasons = ", ".join(set(ssh_failure_reasons)),
            ))
    
    def open_file(self, filepath, mode = "w+"):
        self._ensure_sftp_is_ready()
        sftp_file = self.__sftp_client.file( filepath, mode )
        return _SFTPFileAdapter(sftp_file, filepath)
    
    def open_tempfile(self):
        temp_filename = None
        for stdout_line, stderr_line, stdout_byte, stderr_byte in self.create_command_stream("python -c 'import tempfile; print tempfile.mkstemp()[1]'"):
            if stdout_line:
                temp_filename = stdout_line
        
        sftp_file = self.open_file(temp_filename, "w+")
        return _SFTPFileAdapter(sftp_file, temp_filename)
    
    def create_command_stream(self, command):

        self._ensure_connected_and_authenticated()
        
        # open an SSH shell
        channel = self.__ssh_client.get_transport().open_session()
        
        # get a pseudo-terminal so that certain commands (e.g. scp, rsync, su) behave more nicely
        channel.get_pty()
        
        # grab the streams on that channel
        bufsize = -1
        stdin = channel.makefile('wb', bufsize)
        stdout = channel.makefile('rb', bufsize)
        stderr = channel.makefile_stderr('rb', bufsize)

        # create a callback that will execute the command and return its status
        channel.exec_command(command)
        
        stdout_done_event = Event()
        stderr_done_event = Event()
        
        def watch_stdout():
            
            try:
                # yield each line from the stream
                for byte in bytes_from_stream(stdout):
                    yield byte
                
                # raise exceptions on bad return codes
                log_event(self.__logger.info, "stdout_recv_exit_status")
                code = channel.recv_exit_status()
                if code != 0:
                    log_event(self.__logger.info, "stdout_got_nonzero_retcode", dict(code = code, command = command))
                    raise NonzeroReturnStatusError(code, command)
                log_event(self.__logger.info, "stdout_got_proper_retcode")

            finally:
            
                stdout_done_event.set()

        def watch_stderr():
            
            try:
                # yield each line from the stream
                for byte in bytes_from_stream(stderr):
                    yield byte
                
                # raise exceptions on bad return codes
                log_event(self.__logger.info, "stderr_recv_exit_status")
                code = channel.recv_exit_status()
                if code != 0:
                    log_event(self.__logger.info, "stderr_got_nonzero_retcode", dict(code = code, command = command))
                    raise NonzeroReturnStatusError(code, command)
                log_event(self.__logger.info, "stderr_got_proper_retcode")
            
            finally:
                stderr_done_event.set()
                
        def close_channel_when_stdout_and_stderr_are_done():
            log_event(self.__logger.info, "waiting_to_close_channel")
            stdout_done_event.wait()
            stderr_done_event.wait()
            log_event(self.__logger.info, "closing_channel")
            channel.close()
            log_event(self.__logger.info, "closed_channel")
        
        # be sure to close down the channel when done
        Thread(name = "Thread-close_channel_when_stdout_and_stderr_are_done", target = close_channel_when_stdout_and_stderr_are_done).start()
        
        # set up an input handler
        def input_handler(buffer_str):
            stdin.write( buffer_str )
            
        def kill_handler():
            channel.close()
                
        # FIXME: should also put in stdin, once the IOManager code is in place.  This way you can track what input has gone to the process
        return CommandStream(watch_stdout, watch_stderr, input_handler, kill_handler)

        
    def create_generator_stream(self, generator, virtualenv = None):
        if virtualenv == None:
            python_cmd_format = "python %s"
        else:
            python_cmd_format = "%s/bin/python %s" % virtualenv
        return self._create_generator_stream_with_python_cmd(generator, python_cmd_format)
    
    def _create_generator_stream_with_python_cmd(self, generator, python_cmd_format):
        self._ensure_connected_and_authenticated()
        
        # create the pickled execution for this generator
        pickled_execution = PickledExecution(generator)
        
        # create a temporary file to write to
        temp_code_file = self.open_tempfile()
        temp_code_file.write( pickled_execution.code )
        temp_code_file.close()
        
        # execute it via the shell
        python_cmd = python_cmd_format % temp_code_file.name
        channel = self.__ssh_client.get_transport().open_session()
        channel.get_pty()
        bufsize = -1
        stdin = channel.makefile('wb', bufsize)
        stdout = channel.makefile('rb', bufsize)
        stderr = channel.makefile_stderr('rb', bufsize)
        channel.exec_command(python_cmd)
        
        # set up an input handler
        def input_handler(obj):
            pickled_execution.write_object_to_stdin(obj, stdin)
            
        def kill_handler():
            channel.close()
        
        # do the iteration over the resulting python execution
        return GeneratorStream( lambda: pickled_execution.get_response_iterator(stdout, stderr), input_handler, kill_handler )


class RemoteSuConnection(AbstractConnection):
    """connection object that correponds to a remote SSH connection
    that runs every command su'ed as another user"""
    
    def __init__(self, su_from_connection, username, getpass = None):
        self.__su_from_connection = su_from_connection
        self.__username = username
        self.__getpass = getpass
    
    def open_file(self, filepath, mode = "w+"):
        raise NotImplementedError("'su' functionality is experimental and cannot open files as the current user yet")
    
    def open_tempfile(self, filepath, mode = "w+"):
        raise NotImplementedError("'su' functionality is experimental and cannot open files as the current user yet")
    
    def create_command_stream(self, command):
        if "\"" in command:
            raise ValueError("RemoteSuConnection does not yet support the command '%s'" % command)
        
        # wrap the command and execute it on the su_from RemoteConnection
        su_wrapped_command = "su %s -c \"%s\"" % (self.__username, command)
        command_stream = self.__su_from_connection.create_command_stream(su_wrapped_command)
        
        # wrap the CommandStream we got, by watching stdout/stderr
        # for the first password prompt and responding
        password_prompt_pattern = re.compile(".*assword:.*")
        def password_enabled_stream_iteration():
            pending_stdout_line = ""
            pending_stderr_line = ""
            did_respond_to_password = False
            for stdout_line, stderr_line, stdout_byte, stderr_byte in command_stream:
                
                # track each line as it is built
                if stdout_byte:
                    if stdout_byte == "\n":
                        pending_stdout_line = ""
                    else:
                        pending_stdout_line += stdout_byte
                if stderr_byte:
                    if stderr_byte == "\n":
                        pending_stderr_line = ""
                    else:
                        pending_stderr_line += stderr_byte
                
                # respond to passwords
                if not did_respond_to_password:
                    if password_prompt_pattern.search(pending_stdout_line):
                        password = self.__getpass()
                        command_stream.inject(password + "\r\n")
                        did_respond_to_password = True
                    elif password_prompt_pattern.search(pending_stderr_line):
                        password = self.__getpass()
                        command_stream.inject(password + "\r\n")
                        did_respond_to_password = True
                
                # yield the exact set of data that we *would* have yielded
                yield stdout_line, stderr_line, stdout_byte, stderr_byte
        
        # declare a helper class that will wrap this stream
        class _SuCommandStream(object):
            def __init__(self, command_stream, password_stream_iterator):
                self.__command_stream = command_stream
                self.__password_stream_iterator = password_stream_iterator
            def __iter__(self):
                for stdout_line, stderr_line, stdout_byte, stderr_byte in self.__password_stream_iterator():
                    yield stdout_line, stderr_line, stdout_byte, stderr_byte
            def inject(self, *args, **kwargs):
                self.__command_stream.inject(*args,**kwargs)
            def drain(self, *args, **kwargs):
                self.__command_stream.drain(*args,**kwargs)
            def kill(self, *args, **kwargs):
                self.__command_stream.kill(*args,**kwargs)
                
        # return our wrapped CommandStream
        return _SuCommandStream( command_stream, password_enabled_stream_iteration )
    
    def create_generator_stream(self, generator, virtualenv = None):
        raise NotImplementedError("'su' functionality is experimental and cannot execute remote Python code yet")
    